#' @TODO 预后效能分析，风险因子联动图和KM ROC曲线合图
#' @title ## 预后效能分析，风险因子联动图和KM ROC曲线合图
#' @param model_coef lasso分析结果，第一列为基因名，第二列为系数数值；
#' @param exp 为表达谱
#' @param dataset 为KM_ROC_res结果中的第三个元素，有sample、riskscore、riskgroup，可以为NULL，则默认使用`KM_ROC_res`中的第三个元素
#' @param KM_ROC_res 为之前函数得到的结果，元素一为KM曲线ggsurplot对象，元素二为ROC曲线是ggplot对象
#' 
#' @examples model_coef数据示例
#' > head(lasso_res[[1]])
#'   symbol        coef
#' 1  SPRY2 -0.01705138
#' 2 CCDC69  0.11278595
#' 3 TUBA1C  0.13265976
#' 4  SFRP1 -0.03003526
#' 5 PSMD14  0.11354158
#' 6   RGS2 -0.05530522
#' 
#' @examples exp数据示例
#' > brca_exp_01[1:4,1:4]
#'         TCGA-A2-A0CY-01A TCGA-B6-A40B-01A TCGA-AO-A0J8-01A TCGA-A8-A08J-01A
#' HEPACAM       0.00000000       0.01093413        0.0178391       0.01105015
#' LEP           0.06739103       0.80698082        0.1018269       0.04282505
#' STARD9        0.23719877       0.40060002        0.2003482       0.30404551
#' ANKRD53       0.13371715       0.58099489        0.1788981       0.20353212
#' 
#' @param KM_ROC_res KM_ROC_curve函数分析结果 
#' @examples KM_ROC_res 数据示例
#' > training
#' [[1]]
#' 
#' [[2]]
#' 
#' [[3]]
#'               sample status         time riskgroup  riskscore
#' 1   TCGA-A2-A0CY-01A      0  55.76666667      High -1.7312376
#' 2   TCGA-B6-A40B-01A      0 105.06666667       Low -3.1421083
#' 3   TCGA-AO-A0J8-01A      0  22.66666667       Low -3.8473514
#' @param output_dir 结果输出目录
#' @param var_name 用于命名文件夹,如果为NULL则使用四个随机字符命名
#' @param surtime_unit 生存时间单位 12或者365对应，月份或者天
#' 
#' @return *list*，多张图，均为gg对象
#' @author *WYK*
#'
Effectiveness_Analysis <- function(model_coef = NULL, exp = NULL, clinical = NULL, surtime_unit = c(1,12,365),
                                    saveplot = F, output_dir = NULL, var_name = NULL,dataset = NULL, KM_ROC_res = NULL) { 
  if (is.null(var_name)) {
    var_name <- paste0("_",paste0(sample(letters,4),collapse = ''))
  }

  if(is.null(dataset)){
    dataset <- KM_ROC_res[[3]]
  }

  library(tidyverse)
  library(cowplot)

  p1 <- dataset %>%
    ggplot(mapping = aes(x = reorder(sample, riskscore), y = riskscore)) +
    geom_point(aes(colour = factor(riskgroup, labels = c("High", "Low"))),
      position = "dodge", stat = "identity"
    ) +
    theme_test() +
    labs(x = "Patients(increasing score)", y = "Score") +
    theme(
      text = element_text(size = 12),
      plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"),
      legend.position = "top", axis.ticks.x = element_blank()
    ) +
    scale_x_discrete(labels = NULL) +
    guides(colour = guide_legend(title = "Score")) +
    scale_colour_manual(aesthetics = "colour", values = c("High" = "#E41A1C", "Low" = "#367EB8"))

    if(surtime_unit == 12){
      y_chara <- "Survival Time (Month)"
    }else if (surtime_unit == 365) {
      y_chara <- "Survival Time (Day)"
    }else {
      y_chara <- "Survival Time"
    }

  # scales::show_col(ggsci::pal_igv()(12))

  p2 <- ggplot(data = dataset, aes(x = reorder(sample, riskscore), y = time)) +
    geom_point(aes(colour = factor(status, label = c("Alive", "Dead"))), size = 1.4, alpha = .9) +
    theme_test() +
    labs(x = "Patients(increasing score)", y = y_chara) +
    scale_colour_manual(aesthetics = "colour", values = c("#749B58FF", "#802268FF")) +
    # scale_color_manual(labels = c("Alive", "Dead")) +
    guides(color = guide_legend(title = NULL)) +
    theme(
      text = element_text(size = 12),
      plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"),
      legend.position = "top",
      axis.ticks.x = element_blank()
    ) +
    scale_x_discrete(labels = NULL)

    # gene_index <- which(!is.na(match(rownames(exp),model_coef[, 1])))
    heatmap_data <- na.omit(exp[model_coef[, 1], dataset %>% arrange(riskscore) %>% .$sample])

    group_infor_1 <- dataset %>%
      arrange(riskscore) %>%
      column_to_rownames("sample") %>%
      dplyr::select(riskgroup) %>%
      dplyr::rename(group = 1)

    p3 <- Heatmap_in_Effctiveness(
      data_input = heatmap_data, color_used = c("#E41A1C", "#377EB8"),
      cluster_name = "Score", group_infor = group_infor_1,
      saveplot = F, output_dir = NULL,
      var_name = "Score", width_used = 5, height_used = 4.2, heatmap_name = " "
    )

  # annotation_col <- data.frame(
  #   row.names = dataset %>% arrange(riskscore) %>% pull(sample),
  #   Risk_Group = dataset %>% arrange(riskscore) %>% pull(riskgroup)
  # )

  # ann_colors <- list(
  #   Risk_Group = c(`High` = "#E41A1C", `Low` = "#367EB8")
  # )

 
  # exp <- na.omit(exp)

  # names(heatmap_data) <- names(heatmap_data)

  # p_tmp <- pheatmap::pheatmap(heatmap_data,
  #   scale = "row",
  #   color = colorRampPalette(c("navy", "white", "firebrick3"))(50),
  #   cluster_cols = F,
  #   cluster_rows = F,
  #   show_colnames = F,
  #   show_rownames = T,
  #   # breaks = bk,
  #   annotation_col = annotation_col,
  #   annotation_colors = ann_colors,
  #   legend = T,
  #   annotation_legend = F,
  #   annotation_names_col = FALSE
  #   # filename = '%soutput/p_heatmap.pdf'
  # )

  # p3 <- ggplotify::as.ggplot(p_tmp)


  # p3 <- p3+theme(plot.margin = unit(c(.1,.1,.2,.1),'cm'))

  p1_3 <- plot_grid(p1, p2, p3, ncol = 1, labels = "AUTO", label_size = 20, align = "v", rel_heights = c(.25, .25, .5))

  km <- ggarrange(KM_ROC_res[[1]]$plot + theme(axis.title.x.bottom = element_blank()),
      KM_ROC_res[[1]]$table,
      ncol = 1, align = "v", heights = c(.7, .31)
    )

  # km <- ggpubr::ggarrange(KM_ROC_res[[1]]$plot, KM_ROC_res[[1]]$table, ncol = 1, align = "v", heights = c(0.77, 0.29))

  if (!is.null(KM_ROC_res[[2]])) {
    kmroc_figure <- plot_grid(km, KM_ROC_res[[2]],
      scale = c(1, .965),
      nrow = 2,
      rel_heights = c(1, .97),
      labels = LETTERS[4:5],
      label_size = 20
    )

    library(patchwork)
    lay_out <- "
    111222
    111222
    111222
    111222
    111222
    111222
    "
    p <- p1_3 + kmroc_figure + plot_layout(design = lay_out)

    dir_name <- sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s",output_dir, var_name)
    
    if (saveplot) {

    if(!dir.exists(dir_name)){
      dir.create(dir_name,recursive = T)
    }

      p1_2 <- cowplot::plot_grid(p1, p2,ncol = 1,align = 'v')+theme(plot.margin = unit(c(.2,.2,.2,.2),'cm'))
      p_no_heatroc <- cowplot::plot_grid(km, p1_2,ncol = 2,labels = 'AUTO',align = 'h')

      ggsave2(
        filename = paste0(dir_name, "/Fig_km_scatter.pdf"),
        plot = p_no_heatroc, width = 9.4, height = 5
      )

      ggsave2(
        filename = paste0(dir_name, "/Fig_A-E.pdf"),
        plot = p, width = 10, height = 10
      )
      # ggsave2(
      #   filename = paste0(dir_name, "/Fig_A-E.tiff"),
      #   plot = p, width = 10, height = 10, dpi = 300
      # )
      # ggsave2(
      #   filename = paste0(dir_name, "/Fig_A-E_dpi72.tiff"),
      #   plot = p, width = 10, height = 10, dpi = 72
      # )
      ggsave2(
        filename = paste0(dir_name, "/FigE.pdf"),
        plot = KM_ROC_res[[2]], width = 5, height = 5
      )
    }
  }

  library(cowplot)

  p_scatter <- plot_grid(p1,p2,nrow = 2,align = 'v')
  p_no_heat <- plot_grid(km,p_scatter,align = 'h',nrow = 1,labels = 'AUTO')

  if (saveplot) {
    if (!dir.exists(sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s",output_dir, var_name))) {
      dir.create(sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s",output_dir, var_name),recursive = T)
    } else {
      print(sprintf("Dir '%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s' is existed.",output_dir, var_name))
    }

    dir_name <- sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s",output_dir, var_name)

    ggsave2(
      filename = paste0(dir_name, "/FigA.pdf"),
      plot = p1, width = 4, height = 2.3
    )

    ggsave2(
      filename = paste0(dir_name, "/FigB.pdf"),
      plot = p2, width = 4, height = 2.3
    )

    ggsave2(
      filename = paste0(dir_name, "/FigC.pdf"),
      plot = p3, width = 4, height = 6
    )

    ggsave2(
      filename = paste0(dir_name, "/FigABC.pdf"),
      plot = p1_3, width = 4, height = 9
    )
    
    ggsave2(
      filename = paste0(dir_name, "/FigD.pdf"),
      plot = km, width = 5, height = 5
    )

    ggsave2(
      filename = paste0(dir_name, "/Fig_no_heat.pdf"),
      plot = p_no_heat, width = 9, height = 5
    )
    
  }

  # print(p)

  tmp_list <- list()
  tmp_list[[1]] <- p1
  tmp_list[[2]] <- p2
  tmp_list[[3]] <- p3
  tmp_list[[4]] <- p1_3
  tmp_list[[5]] <- p

  names(tmp_list) <- c("p1", "p2", "p3", "risk_figure", "1_5p")
  return(tmp_list)
}
# Effectiveness_Analysis的子函数，用于绘制热图
Heatmap_in_Effctiveness <- function(data_input = NULL, color_used= NULL, cluster_name= "cluster", group_infor = NULL,
                          saveplot = T, output_dir= './', var_name= NULL, width_used = 9, height_used = 9, heatmap_name = " ") {

    if(is.null(var_name)){
        var_name <- paste0("_",paste0(sample(letters,4), collapse = '',sep = ""))
    }
    group_infor_2 <- group_infor

    library(tidyverse)
    library(ComplexHeatmap)
    library(cowplot)
    library(ggpubr)

    sample_common <- intersect(rownames(group_infor),colnames(data_input))
    group_infor <- group_infor[sample_common,,drop = F]
    data_input <- data_input[,sample_common]
    
    group_chara <- sort(unique(as.character(group_infor$group)))

    if (is.null(color_used)) {
        col_name <- RColorBrewer::brewer.pal(8, "Set1")[1:length(group_chara)]
        col_name <- col_name[seq_along(unique(group_infor %>% pull(group)))]
    } else {
        col_name <- color_used[seq_along(unique(group_infor %>% pull(group)))]
    }

    names(col_name) <- group_chara

    if (ncol(group_infor) >= 2) {
        col_anno <- HeatmapAnnotation(
            df = group_infor, # %>% .[, "group", drop = F]
            col = list(
                group = col_name
            ),
            annotation_legend_param = list(group = list(title = c(cluster_name, colnames(group_infor)[2:ncol(group_infor)]))),
            annotation_name_side = "left",
            annotation_label = c(cluster_name, colnames(group_infor)[2:ncol(group_infor)])
        )
    } else {
        col_anno <- HeatmapAnnotation(
            df = group_infor, # %>% .[, "group", drop = F]
            col = list(
                group = col_name
            ),
            annotation_name_side = "left",
            annotation_label = cluster_name,
            annotation_legend_param = list(group = list(title = cluster_name))
        )
    }

    data_input_scaled <- t(apply(data_input, 1, function(x) scale(x, center = T, scale = T)))
    colnames(data_input_scaled) <- colnames(data_input)
    all(colnames(data_input_scaled) == group_infor %>% rownames())
    

    group_infor_2 <- group_infor_2[sample_common,,drop = F]
    col_zscore <- circlize::colorRamp2(c(-4, 0, 4), c("navy", "white", "firebrick3")) # c("#0a5aa5", "white", "firebrick3")

    cluster_res <- Heatmap(
        matrix = data_input_scaled,
        name = heatmap_name,
        col = col_zscore,
        show_column_dend = F,
        show_column_names = F,
        show_row_names = T,
        # cluster_rows = T,
        show_row_dend = F,
        # rect_gp = gpar(col = "white", lwd = 1),
        row_order = rownames(data_input_scaled),
        row_names_gp = gpar(fontsize = 9.5),
        column_title = " ",
        top_annotation = col_anno,
        column_order = rownames(group_infor_2),
        # heatmap_legend_param = list(direction = "horizontal"),
        # column_split = group_infor %>% .[, 1, drop = F],
        border = F
        # left_annotation = row_anno,
        # row_names_side = "right"
    )

    gg.cluster_res <- cluster_res %>%
        draw(., merge_legend = T,heatmap_legend_side = "left", 
                annotation_legend_side = "left") %>%
        grid.grabExpr() %>%
        ggplotify::as.ggplot() +
        theme(
            plot.background = element_rect(fill = "white", color = "white"),
            plot.margin = unit(c(.1, .1, .3, .1), "cm")
        )

    if (saveplot) {
        dir_now <- str_glue("{output_dir}/output/heatmap/")

        if (!dir.exists(dir_now)) {
            dir.create(dir_now, recursive = T)
        } else {
            print(str_c(dir_now, " is ready."))
        }

        ggsave2(
            filename = str_glue("{dir_now}/figure_heat_plot{var_name}.pdf"),
            plot = gg.cluster_res, width = width_used, height = height_used
        )
        ggsave2(
            filename = str_glue("{dir_now}/figure_heat_plot{var_name}.tiff"),
            plot = gg.cluster_res, width = width_used, height = height_used, dpi = 300
        )
        ggsave2(
            filename = str_glue("{dir_now}/figure_heat_plot{var_name}_dpi72.tiff"),
            plot = gg.cluster_res, width = width_used, height = height_used, dpi = 72
        )
    }
    return(gg.cluster_res)
}

